#include "include/core/SkCanvas.h"
#include "include/core/SkBitmap.h"
#include "dsp/DCT.h"
#include "util/UUID.h"
#include <array>
#include <cmath>
#include <vector>
#include <algorithm>

using namespace yzrilyzr_dsp;
using namespace yzrilyzr_util;
using namespace yzrilyzr_lang;

static DCT dct8x8(8); // 8x8 DCT变换

// 生成校验码 (CRC8简化版)
static uint8_t generateChecksum(const std::vector<bool> & bits){
	uint8_t crc=0;
	for(bool bit:bits){
		crc^=(bit << (crc % 8));
		crc=(crc << 1) | (crc >> 7);
	}
	return crc;
}

// UUID转换为比特流 (128bit UUID + 8bit校验码)
static std::vector<bool> uuidToBits(uint64_t msb, uint64_t lsb){
	std::vector<bool> bits;
	bits.reserve(136); // 128 + 8

	// 添加UUID比特
	for(u_index i=63; i >= 0; --i) bits.push_back((msb >> i) & 1);
	for(u_index i=63; i >= 0; --i) bits.push_back((lsb >> i) & 1);

	// 添加校验码
	uint8_t checksum=generateChecksum(bits);
	for(u_index i=7; i >= 0; --i) bits.push_back((checksum >> i) & 1);

	return bits;
}

// 嵌入盲水印（优化版）
static void embedBlindWatermark(SkCanvas & canvas, const UUID & uuid){
	SkBitmap bitmap;
	if(!bitmap.tryAllocN32Pixels(canvas.imageInfo().width(), canvas.imageInfo().height())){
		return;
	}
	if(!canvas.readPixels(bitmap, 0, 0)){
		return;
	}

	const int width=bitmap.width();
	const int height=bitmap.height();
	uint32_t * pixels=reinterpret_cast<uint32_t *>(bitmap.getPixels());

	// 1. 提取Y通道（亮度）
	std::vector<float> yChannel(width * height);
	for(u_index i=0; i < width * height; ++i){
		uint32_t pixel=pixels[i];
		uint8_t r=(pixel >> 16) & 0xFF;
		uint8_t g=(pixel >> 8) & 0xFF;
		uint8_t b=pixel & 0xFF;
		yChannel[i]=0.299f * r + 0.587f * g + 0.114f * b;
	}

	// 2. 嵌入水印
	const std::vector<bool> bits=uuidToBits(uuid.msb, uuid.lsb);
	const int bitSize=bits.size();
	int bitIndex=0;

	// 计算每个8x8块需要嵌入的bit数（全画布填充）
	const int blocksX=(width + 7) / 8;
	const int blocksY=(height + 7) / 8;
	const int bitsPerBlock=(bitSize + blocksX * blocksY - 1) / (blocksX * blocksY);

	for(int i=0; i < height; i+=8){
		const int blockHeight=Math::min(8, height - i);
		for(int j=0; j < width; j+=8){
			const int blockWidth=Math::min(8, width - j);

			// 2.1 提取8x8块
			std::array<float, 64> block={};
			for(u_index y=0; y < blockHeight; ++y){
				for(u_index x=0; x < blockWidth; ++x){
					block[y * 8 + x]=yChannel[(i + y) * width + (j + x)];
				}
			}

			// 2.2 DCT变换
			std::array<float, 64> dctBlock={};
			dct8x8.dct(block.data(), dctBlock.data());

			// 2.3 在当前块嵌入多个bit（提高鲁棒性）
			for(u_index k=0; k < bitsPerBlock && bitIndex < bitSize; ++k){
				// 选择多个中频系数嵌入（提高抗压缩能力）
				const int pos1=3 * 8 + 3; // 中频区域
				const int pos2=4 * 8 + 4;
				const int pos3=5 * 8 + 3;

				// 嵌入强度自适应（基于块亮度）
				float avgLuma=0;
				for(u_index n=0; n < 64; ++n) avgLuma+=block[n];
				avgLuma/=64;
				const float strength=5.0f + avgLuma * 0.05f; // 5-18强度范围

				// 修改DCT系数
				if(bits[bitIndex++]){
					dctBlock[pos1]+=strength;
					dctBlock[pos2]+=strength * 0.8f;
					dctBlock[pos3]+=strength * 0.6f;
				} else{
					dctBlock[pos1]-=strength;
					dctBlock[pos2]-=strength * 0.8f;
					dctBlock[pos3]-=strength * 0.6f;
				}
			}

			// 2.4 IDCT逆变换
			std::array<float, 64> idctBlock={};
			dct8x8.idct(dctBlock.data(), idctBlock.data());

			// 2.5 写回Y通道
			for(u_index y=0; y < blockHeight; ++y){
				for(u_index x=0; x < blockWidth; ++x){
					yChannel[(i + y) * width + (j + x)]=idctBlock[y * 8 + x];
				}
			}
		}
	}

	// 3. 更新RGB数据（YUV→RGB）
	for(u_index i=0; i < width * height; ++i){
		const float Y=yChannel[i];
		const uint32_t pixel=pixels[i];
		const uint8_t r=(pixel >> 16) & 0xFF;
		const uint8_t g=(pixel >> 8) & 0xFF;
		const uint8_t b=pixel & 0xFF;

		// 保持色度不变
		const float U=-0.147f * r - 0.289f * g + 0.436f * b;
		const float V=0.615f * r - 0.515f * g - 0.100f * b;
		const uint8_t newR=static_cast<uint8_t>(std::clamp(Y + 1.140f * V, 0.0f, 255.0f));
		const uint8_t newG=static_cast<uint8_t>(std::clamp(Y - 0.395f * U - 0.581f * V, 0.0f, 255.0f));
		const uint8_t newB=static_cast<uint8_t>(std::clamp(Y + 2.032f * U, 0.0f, 255.0f));

		pixels[i]=(0xFF << 24) | (newR << 16) | (newG << 8) | newB;
	}

	// 4. 写回画布
	canvas.writePixels(bitmap.info(), pixels, width * sizeof(uint32_t), 0, 0);
}